{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-10-29T13:08:16.954658Z",
     "start_time": "2025-10-29T13:07:10.486704Z"
    }
   },
   "source": [
    "# mnist_pytorch.py\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "# ========== 1. 数据加载 ==========\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),                     # 转为 [0,1]\n",
    "    transforms.Normalize((0.1307,), (0.3081,)) # 标准化 MNIST 均值/方差\n",
    "])\n",
    "\n",
    "train_dataset = datasets.MNIST(root=\"./data\", train=True, transform=transform, download=True)\n",
    "test_dataset  = datasets.MNIST(root=\"./data\", train=False, transform=transform, download=True)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "test_loader  = DataLoader(test_dataset, batch_size=1000, shuffle=False)\n",
    "\n",
    "# ========== 2. 定义 CNN 模型 ==========\n",
    "class CNNModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(CNNModel, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, 3, 1)\n",
    "        self.conv2 = nn.Conv2d(32, 64, 3, 1)\n",
    "        self.dropout1 = nn.Dropout(0.25)\n",
    "        self.dropout2 = nn.Dropout(0.5)\n",
    "        self.fc1 = nn.Linear(9216, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.conv1(x))\n",
    "        x = F.relu(self.conv2(x))\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = self.dropout1(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.dropout2(x)\n",
    "        x = self.fc2(x)\n",
    "        output = F.log_softmax(x, dim=1)\n",
    "        return output\n",
    "\n",
    "# ========== 3. 初始化 ==========\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = CNNModel().to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# ========== 4. 训练与验证 ==========\n",
    "def train(model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "    print(f\"Epoch {epoch}: Train loss = {total_loss / len(train_loader):.4f}\")\n",
    "\n",
    "def test(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            test_loss += F.nll_loss(output, target, reduction='sum').item()\n",
    "            pred = output.argmax(dim=1)\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    acc = correct / len(test_loader.dataset)\n",
    "    print(f\"Test set: Average loss = {test_loss:.4f}, Accuracy = {acc*100:.2f}%\")\n",
    "    return acc\n",
    "\n",
    "# ========== 5. 主训练循环 ==========\n",
    "epochs = 5\n",
    "for epoch in range(1, epochs + 1):\n",
    "    train(model, device, train_loader, optimizer, epoch)\n",
    "    test(model, device, test_loader)\n",
    "\n",
    "# ========== 6. 保存模型 ==========\n",
    "torch.save(model.state_dict(), \"mnist_cnn.pt\")\n",
    "print(\"✅ 模型已保存为 mnist_cnn.pt\")\n",
    "\n",
    "# ========== 7. 单张图像预测 ==========\n",
    "def predict_image(img_path, model_path=\"mnist_cnn.pt\"):\n",
    "    model = CNNModel()\n",
    "    model.load_state_dict(torch.load(model_path, map_location='cpu'))\n",
    "    model.eval()\n",
    "\n",
    "    # 读入灰度图像\n",
    "    img = Image.open(img_path).convert('L')\n",
    "    img = img.resize((28, 28))\n",
    "    arr = np.array(img).astype(np.float32)\n",
    "\n",
    "    # 反转颜色（白底黑字 → 黑底白字）\n",
    "    if arr.mean() > 127:\n",
    "        arr = 255 - arr\n",
    "\n",
    "    arr = arr / 255.0\n",
    "    arr = (arr - 0.1307) / 0.3081  # 同样的标准化\n",
    "    tensor = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0)  # (1,1,28,28)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        output = model(tensor)\n",
    "        pred = output.argmax(dim=1).item()\n",
    "        confidence = torch.exp(output[0, pred]).item()\n",
    "\n",
    "    plt.imshow(arr, cmap='gray')\n",
    "    plt.title(f\"Predicted: {pred} (conf={confidence:.3f})\")\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "    return pred, confidence\n",
    "\n",
    "# 示例（替换为你自己的图片路径）\n",
    "# predict_image(\"my_digit.png\")\n"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: Train loss = 0.2062\n",
      "Test set: Average loss = 0.0441, Accuracy = 98.42%\n",
      "Epoch 2: Train loss = 0.0855\n",
      "Test set: Average loss = 0.0384, Accuracy = 98.78%\n",
      "Epoch 3: Train loss = 0.0649\n",
      "Test set: Average loss = 0.0324, Accuracy = 99.00%\n",
      "Epoch 4: Train loss = 0.0535\n",
      "Test set: Average loss = 0.0298, Accuracy = 99.00%\n",
      "Epoch 5: Train loss = 0.0455\n",
      "Test set: Average loss = 0.0314, Accuracy = 99.06%\n",
      "✅ 模型已保存为 mnist_cnn.pt\n"
     ]
    }
   ],
   "execution_count": 1
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
